from mpl_toolkits import mplot3d
import numpy as np
import matplotlib.pyplot as plt
from python_ai.ML.lin_regression.xlib import *


def func(X):
    # return 2.5 * X[:, 0] + 0.7 * X[:, 1] + 1.2 * X[:, 2]
    return np.dot(X, np.array([2.5, 0.7, 1.2]).T)


np.random.seed(1)
m = 20

XarrAx = np.linspace(1, m, m)
YarrAx = XarrAx.copy()
print(YarrAx)
Yarr = YarrAx + np.random.normal(0, 1, m) * 0.5
print(Yarr)
X = np.c_[np.ones([m, 1]), XarrAx.reshape(m, 1), Yarr.reshape(m, 1)]
Y = (func(X) + np.random.normal(0, 1, m) * 0.5).reshape(m, 1)
print(X.shape, Y.shape)

num_iters = 100
theata, history, xscores = gradient_descent_algorithm(X, Y, num_iters=num_iters)
print(theata)


def func_fit(X, theata):
    Y = theata[0] * X[:, 0] + theata[1] * X[:, 1] + theata[2] * X[:, 2]
    return Y


def on_ax(ax):
    ax.set_title('梯度下降算法')
    ax.scatter3D(XarrAx, Yarr, Y)
    Xy = np.c_[np.ones([m, 1]), XarrAx.reshape(m, 1), YarrAx.reshape(m, 1)]
    ax.plot3D(XarrAx, YarrAx, func_fit(Xy, theata), 'r-')
    ax.plot3D(XarrAx, YarrAx, func(Xy), 'g--')
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')


plt.ioff()

fig = plt.figure(figsize=[12, 8])
pr = 1
pc = 1
plt.rcParams['font.sans-serif'] = ['Simhei']  # 设置中文字体为黑体
plt.rcParams['axes.unicode_minus'] = False  # 显示负号

ax = fig.add_subplot(pr, pc, 1, projection='3d')
on_ax(ax)

fig = plt.figure(figsize=[12, 8])
pr = 2
pc = 2
plt.rcParams['font.sans-serif'] = ['Simhei']  # 设置中文字体为黑体
plt.rcParams['axes.unicode_minus'] = False  # 显示负号

ax = fig.add_subplot(pr, pc, 1, projection='3d')
on_ax(ax)

plt.subplot(pr, pc, 2)
plt.title('Cost function value')
xx = range(num_iters)
plt.plot(xx, history)
plt.grid()

plt.subplot(pr, pc, 3)
plt.title('Cost function value, 2nd half')
xx = range(num_iters // 2, num_iters)
plt.plot(xx, history[num_iters // 2:])
plt.grid()

plt.subplot(pr, pc, 4)
plt.title('Score function value')
xx = range(num_iters)
plt.plot(xx, xscores)
plt.grid()

plt.show()
